#9月26日改
#避免联网下载大模型
import os
os.environ["HF_HUB_OFFLINE"] = "1"#禁止任何 HTTP 请求
os.environ["TRANSFORMERS_OFFLINE"] = "1"#让 transformers 也不联网
os.environ["HF_HOME"] = "./model"  # 把缓存目录指到项目内，减少磁盘寻址时间，避免系统临时目录杀毒扫描

###
## Web 框架核心
from flask import (Flask, render_template, request, redirect,url_for, flash, jsonify, Response)
# ORM，用于操作 SQLite
from flask_sqlalchemy import SQLAlchemy
## 用户登录/会话管理
from flask_login import (LoginManager, login_user, logout_user,login_required, current_user, UserMixin)
# 密码加解密
from werkzeug.security import generate_password_hash, check_password_hash
# 调用 LLM
from openai import OpenAI
# 自行写的配置文件：放 key 和 base_url
from config import API_CONFIG  # 我的 key & base_url
##引入日期时间库
import  datetime as dt
# 把系统时区强制设为上海
os.environ['TZ'] = 'Asia/Shanghai'
# 引入 tempfile，为打印临时目录
import  tempfile

print("cwd:", os.getcwd())#   # 打印当前工作目录
print("temp:", tempfile.gettempdir())# # 打印系统临时目录


#读取 flask 的密钥
from config import SECRET_KEY
# 读取要使用的具体模型名,.env的优先级高于config
from config import DEEPSEEK_MODEL

MODEL = DEEPSEEK_MODEL
print(">>> 当前 LLM 模型:", MODEL)

#再引一次，供后续 Tavily 搜索用
import os, requests, json

# 向量化 - 使用正确的导入方式
#from langchain_community.embeddings import HuggingFaceEmbeddings
#10月11日改，消除LangChainDeprecationWarning
from langchain_huggingface import HuggingFaceEmbeddings## 本地向量嵌入
from langchain_chroma import Chroma# 向量数据库
from langchain.text_splitter import RecursiveCharacterTextSplitter#文本切块
from langchain_community.document_loaders import TextLoader, PyPDFLoader#各类文件加载器
from langchain_community.document_loaders import UnstructuredWordDocumentLoader, UnstructuredPowerPointLoader, CSVLoader


#解决Tavily 超时（网络问题）
from langchain_community.tools import DuckDuckGoSearchRun## 搜索兜底方案
ddg = DuckDuckGoSearchRun()#实例化


# 注册表
# 文件扩展名 → 对应加载器类
LOADER_MAP = {
    ".txt": TextLoader,
    ".md": TextLoader,
    ".csv": CSVLoader,
    ".pdf": PyPDFLoader,  # 文字型 PDF
    ".docx": UnstructuredWordDocumentLoader,
    ".pptx": UnstructuredPowerPointLoader,
    ".html": TextLoader,
    ".htm":  TextLoader,
}


# 自定义加载器选择函数
def load_document(file_path: str):    #根据后缀自动选加载器并读文件
    ext = os.path.splitext(file_path)[1].lower() # 取小写后缀
    loader_class = LOADER_MAP.get(ext) # 查表
    if not loader_class: # 不支持的格式返回
        return None

    try:
        if loader_class == TextLoader:# TextLoader 需要额外指定编码
            loader = loader_class(file_path, encoding="utf-8")
        else:
            loader = loader_class(file_path)
        return loader.load()# 读取并返回 Document 列表
    except Exception as e:
        print(f"加载文件 {file_path} 时出错: {e}")
        return None


# 修改后的向量库构建函数
def build_or_load_vector_store():
    persist_dir = "knowledge_db"# 向量库存放路径

    # 使用正确的嵌入模型
    # embeddings = HuggingFaceEmbeddings(
    #     model_name="all-MiniLM-L6-v2"
    # )
    #9月28日改
    # 强制读本地，不再联网
    embeddings = HuggingFaceEmbeddings(
        model_name="./model/all-MiniLM-L6-v2"
    )
    ###
    if os.path.exists(persist_dir):  # 若已存在则直接加载
        print("加载现有的向量数据库...")
        try:
            return Chroma(persist_directory=persist_dir, embedding_function=embeddings)
        except Exception as e:
            print(f"加载现有向量库失败: {e}，重新构建...")
            # 如果加载失败，删除目录重新构建
            import shutil
            shutil.rmtree(persist_dir)

    print("构建新的向量数据库...")

    # 手动加载文档
    docs = []
    my_docs_dir = "my_docs" # 用户上传的文档目

    if not os.path.exists(my_docs_dir): # 目录不存在就建一个
        print(f"警告: {my_docs_dir} 目录不存在，创建空目录")
        os.makedirs(my_docs_dir, exist_ok=True)
        # 创建空向量库
        vectorstore = Chroma.from_documents([], embeddings, persist_directory=persist_dir)
        print("空向量数据库创建完成")
        return vectorstore

    # 遍历目录并加载所有支持的文件
    for root, _, files in os.walk(my_docs_dir):
        for file in files:
            file_path = os.path.join(root, file)
            file_ext = os.path.splitext(file)[1].lower()

            if file_ext in LOADER_MAP:# 只处理支持的格式
                print(f"正在加载: {file_path}")
                loaded_docs = load_document(file_path)
                if loaded_docs:
                    docs.extend(loaded_docs)
            else:
                print(f"跳过不支持的文件格式: {file_path}")

    if not docs: # 处理完仍没文档
        print("未找到可处理的文档，创建空向量库")
        vectorstore = Chroma.from_documents([], embeddings, persist_directory=persist_dir)
        print("空向量数据库创建完成")
        return vectorstore

    # 分割文档
    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
    chunks = splitter.split_documents(docs)

    print(f"共处理 {len(docs)} 个文档，分割为 {len(chunks)} 个片段")

    # 创建向量库 - 新版本 Chroma 会自动持久化
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=persist_dir
    )
    print("向量数据库构建完成并已自动持久化")

    return vectorstore


# 尝试构建向量库，如果失败则创建空库
try:
    vector_store = build_or_load_vector_store() # 全局变量，供后面检索
    print("向量库加载成功")
except Exception as e:
    print(f"向量库加载失败: {e}")
    # 创建空的嵌入模型作为后备

    #embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    #9月26日改，缩短模型启动加载时间
    # embeddings = HuggingFaceEmbeddings(
    #     model_name="./model/all-MiniLM-L6-v2",  # 模型的目录
    #     cache_folder=None,  # 禁用额外缓存
    #     #local_files_only=True  # 强制只读本地，不再访问 HuggingFace
    # )
    # 下面复用本地模型目录，不再访问 HuggingFace
    MODEL_DIR = os.getenv("MODEL_PATH") or os.path.join(
        os.path.abspath(os.path.dirname(__file__)),
        "model",
        "all-MiniLM-L6-v2"
    )

    embeddings = HuggingFaceEmbeddings(model_name=MODEL_DIR, cache_folder=None)
    ###
    vector_store = None# 标记为“空库”
    print("使用空向量库作为后备")


# 改动,适配Tavily api的接口
# def search(query_list, num=6):
#     api_key = os.getenv("TAVILY_API_KEY")
#     if not api_key:
#         print("警告: 未设置 TAVILY_API_KEY 环境变量")
#         return {"error": "API密钥未配置", "results": []}
#
#     url = "https://api.tavily.com/search"
#     payload = {
#         "query": " ".join(query_list),
#         "max_results": num,
#         "search_depth": "basic",
#         "include_answer": False,
#         "include_raw_content": True
#     }
#     headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
#     try:
#         r = requests.post(url, json=payload, headers=headers, timeout=8)
#         r.raise_for_status()
#         return r.json()
#     except Exception as e:
#         print("Tavily 搜索异常:", e)
#         return {"error": str(e), "results": []}
 # 新实现：先 Tavily，失败用 DuckDuckGo
def search(query_list, num=6):
    api_key = os.getenv("TAVILY_API_KEY")
    if not api_key:
        print("警告：未设置 TAVILY_API_KEY，直接回退到 DuckDuckGo")
        return ddg_fallback(" ".join(query_list), num)   # 别写成 ddg.fallback

    url = "https://api.tavily.com/search"
    payload = {
        "query": " ".join(query_list),
        "max_results": num,
        "search_depth": "basic",
        "include_answer": False,
        "include_raw_content": True
    }
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    try:
        r = requests.post(url, json=payload, headers=headers, timeout=8)
        r.raise_for_status()
        return r.json()
    except Exception as e:
        print("Tavily 搜索异常:", e)
        return ddg_fallback(" ".join(query_list), num)


 # DuckDuckGo 兜底，返回 Tavily 格式
def ddg_fallback(query, num):
        """用 DuckDuckGo 兜底，返回和 Tavily 一样的格式"""
        try:
            results = ddg.run(query)  # 字符串结果
            # 简单拆成 num 条
            snippets = [r.strip() for r in results.split("\n") if r.strip()][:num]
            return {
                "results": [
                    {"title": f"结果{i + 1}", "url": "", "content": s, "snippet": s}
                    for i, s in enumerate(snippets)
                ]
            }
        except Exception as e:
            print("DuckDuckGo 也挂了:", e)
            return {"error": str(e), "results": []}

#把搜索返回打包成模型易读的字符串
def format_search_results(results):
    """格式化搜索结果供模型使用"""
    if not results or "error" in results:
        return "未能获取实时搜索结果"

    if "results" not in results or not results["results"]:
        return "未找到相关搜索结果"

    formatted = []
    for item in results["results"][:4]:  # 最多取 4 条
        title = item.get("title", "无标题")
        url = item.get("url", "")
        content = item.get("content", "") or item.get("snippet", "")
        formatted.append(f"标题：{title}\n链接：{url}\n摘要：{content[:500]}…")
    return "\n\n".join(formatted)


# ------------------ 基础配置 ------------------
app = Flask(__name__)# 创建 Flask 实例
app.config['SECRET_KEY'] = SECRET_KEY# 设置 flask 密钥

#改
#10月30日，中文就不会再被 Unicode 转义了
# app.config['JSON_AS_ASCII'] = False
# app.config['JSONIFY_PRETTYPRINT_REGULAR'] = False
# app.config['JSON_SORT_KEYS'] = False
# # 放在 try/except 后面
# print('>>> JSON_AS_ASCII 设置值:', app.config.get('JSON_AS_ASCII'))
# 数据库放在当前目录 chat.db
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + os.path.join(
    os.path.abspath(os.path.dirname(__file__)), 'chat.db'
)

db = SQLAlchemy(app)  # 初始化 ORM
login_manager = LoginManager()
login_manager.init_app(app)# 登录管理绑定 flask
login_manager.login_view = 'login'# 未登录自动跳转到 login 路由

client = OpenAI(api_key=API_CONFIG["api_key"],
                base_url=API_CONFIG["base_url"])# 实例化 LLM 客户端
#MODEL = 'deepseek-reasoner' # 真正使用的模型名


# ------------------ 数据模型 ------------------
class User(UserMixin, db.Model):# 用户表
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)
    is_admin = db.Column(db.Boolean, default=False) # 管理员标志

    def set_password(self, pwd): # 把明文密码哈希后存库
        self.password_hash = generate_password_hash(pwd)

    def check_password(self, pwd): # 验证明文 vs 哈希
        return check_password_hash(self.password_hash, pwd)


class Message(db.Model):# 留言表
    id = db.Column(db.Integer, primary_key=True)
    body = db.Column(db.Text, nullable=False)# 用户留言内容
    author_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    author = db.relationship('User', backref='messages')
    created = db.Column(
        db.DateTime,# 东八区时间
        default=lambda: dt.datetime.now(dt.timezone(dt.timedelta(hours=8)))
    )
    ai_reply = db.Column(db.Text) # AI 回复内容


# ------------------ 登录回调 ------------------
@login_manager.user_loader# flask-login 要求实现的回调
def load_user(user_id):
    return db.session.get(User, int(user_id))# SQLAlchemy 2.0 语法


# ------------------ 路由：认证 ------------------
@app.route('/register', methods=['GET', 'POST'])# 注册页面
def register():
    if request.method == 'POST':
        username = request.form['username'].strip()
        pwd = request.form['password']
        if User.query.filter_by(username=username).first():
            flash('用户名已存在')
            return redirect(url_for('register'))
        u = User(username=username)
        u.set_password(pwd)
        db.session.add(u)
        db.session.commit()
        flash('注册成功，请登录')
        return redirect(url_for('login'))
    return render_template('register.html')


@app.route('/login', methods=['GET', 'POST'])# 登录页面
def login():
    if request.method == 'POST':
        u = User.query.filter_by(username=request.form['username']).first()
        if u and u.check_password(request.form['password']):
            login_user(u)
            return redirect(url_for('index'))
        flash('用户名或密码错误')
    return render_template('login.html')


@app.route('/logout')
@login_required # 必须登录才能访问
def logout():
    logout_user()
    return redirect(url_for('login'))


# ------------------ 路由：留言板 ------------------
@app.route('/')
@login_required
def index():
    msgs = Message.query.order_by(Message.created.desc()).all()
    return render_template('index.html', msgs=msgs)


@app.route('/add', methods=['POST'])
@login_required
def add_message():
    body = request.form['body'].strip()
    if body:
        m = Message(body=body, author=current_user)
        db.session.add(m)
        db.session.commit()
    return redirect(url_for('index'))


# ------------------ 路由：管理员 ------------------
@app.route('/reply/<int:msg_id>')  # SSE 流式返回 AI
def reply_stream(msg_id):
    msg = Message.query.get_or_404(msg_id)

    if msg.ai_reply: # 若已有缓存直接返回
        return jsonify({'content': msg.ai_reply})

    #embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    # 9月26日改
    # embeddings = HuggingFaceEmbeddings(
    #     model_name="./model/all-MiniLM-L6-v2",  # 你放模型的目录
    #     cache_folder=None,  # 禁用额外缓存
    #     #local_files_only=True  # 强制只读本地，不再访问 HuggingFace
    # )


    MODEL_DIR = os.getenv("MODEL_PATH") or os.path.join(
        os.path.abspath(os.path.dirname(__file__)),
        "model",
        "all-MiniLM-L6-v2"
    )

    embeddings = HuggingFaceEmbeddings(model_name=MODEL_DIR, cache_folder=None)
    ###

    if vector_store: # 优先用全局已加载的向量库
        current_vector_store = vector_store
    else:
        try:  # 否则建一个空的
            #current_vector_store = Chroma.from_documents([], embeddings, persist_directory="knowledge_db")
            #9月26日日志改
            current_vector_store = Chroma(
                embedding_function=embeddings,
                persist_directory="knowledge_db",
                collection_metadata={"hnsw:space": "cosine"}
            )

        except Exception as e:
            print(f"创建空向量库失败: {e}")
            return jsonify({'content': '知识库初始化失败，请联系管理员'})

    def generate(): # SSE 生成器
        full = ''
        try:
            if current_vector_store:
                docs = current_vector_store.similarity_search(msg.body, k=20)
                print(f"召回 {len(docs)} 个相关片段")
                private_ctx = "\n".join([d.page_content for d in docs])
            else:
                private_ctx = "知识库暂不可用"

            web_ctx = format_search_results(search([msg.body]))# 实时搜索

            user_prompt = f"【公司知识】{private_ctx}\n【实时搜索】{web_ctx}\n用户问题：{msg.body}"

            stream = client.chat.completions.create(
                model=MODEL,
                messages=[
                    {"role": "system", "content": _make_system_prompt()},
                    {"role": "user", "content": user_prompt}
                ],
                stream=True,
                temperature=0.5,
                max_tokens=2048,
                top_p=0.95
            )

            for chunk in stream:
                delta = chunk.choices[0].delta.content or ""
                if delta:
                    full += delta
                    yield f"data: {json.dumps({'content': delta})}\n\n"

            with app.app_context():# 流式写完后把完整回复存库
                msg.ai_reply = full
                db.session.commit()
        except Exception as e:
            error_msg = f"生成回复时出错: {str(e)}"
            print(error_msg)
            yield f"data: {json.dumps({'content': error_msg, 'error': True})}\n\n"

    return Response(generate(), mimetype='text/event-stream')


@app.route('/admin') # 管理员面板
@login_required
def admin():
    if not current_user.is_admin:
        flash('无权限')
        return redirect(url_for('index'))
    users = User.query.all()
    return render_template('admin.html', users=users)


@app.route('/admin/toggle/<int:user_id>')  # 切换管理员权限
@login_required
def toggle_admin(user_id):
    if not current_user.is_admin:
        return 'Forbidden', 403
    u = User.query.get_or_404(user_id)
    u.is_admin = not u.is_admin
    db.session.commit()
    return redirect(url_for('admin'))

#改
#新增路由
# ------------------ 路由：文件上传 ------------------
@app.route('/upload', methods=['POST'])
@login_required
def upload_file():
    from werkzeug.utils import secure_filename
    file = request.files.get('file')
    if not file or file.filename == '':
        return jsonify({'status': 'error', 'msg': '未选择文件'}), 400
    print("原始文件名:", file.filename)
    filename = secure_filename(file.filename)
    print("secure之后:", filename)
    ext = os.path.splitext(filename)[1].lower()
    if ext not in LOADER_MAP:
        return jsonify({'status': 'error', 'msg': f'不支持的文件类型: {ext}'}), 400
    os.makedirs("my_docs", exist_ok=True)
    save_path = os.path.join("my_docs", filename)
    file.save(save_path)
    global vector_store
    try:
        vector_store = build_or_load_vector_store()
    except Exception as e:
        return jsonify({'status': 'error', 'msg': f'知识库更新失败: {str(e)}'}), 500
    return jsonify({'status': 'ok', 'msg': '上传并更新成功'})


###
# ------------------ 工具函数 ------------------
 # 构造系统提示，注入当前日期
def _make_system_prompt():
    from datetime import datetime, timedelta, timezone
    today = datetime.now(timezone(timedelta(hours=8))).date()
    weekday = ['一', '二', '三', '四', '五', '六', '日'][today.weekday()]
    return f"""
You are DeepSeek Chat, a large language model trained by DeepSeek, based on the DeepSeek-V3.1 architecture.
Knowledge cutoff: 2024-07
今天是 {today.year}年{today.month}月{today.day}日，星期{weekday}。请始终用这句话回答任何关于当前日期的问题。
You can use tools like web search, drawing, code interpreter, etc., when needed.
Answer in the same language as the user.
You are a helpful assistant. Always answer accurately and concisely.

"""


# 配置检查函数
def check_configurations(): # 启动前检查环境变量/目录
    required_env_vars = ["SECRET_KEY"]
    missing_vars = [var for var in required_env_vars if not os.getenv(var)]

    if missing_vars:
        print(f"警告: 缺少环境变量: {', '.join(missing_vars)}")

    if not os.path.exists("my_docs"):
        print("警告: my_docs 目录不存在，知识库功能将不可用")

    if not os.path.exists("knowledge_db"):
        print("信息: 知识向量库不存在，将在首次请求时创建")

#添加接口
#1种类型
# @app.route('/ai/analyze', methods=['POST'])
# def ai_analyze():
#     data = request.get_json()
#     demand_content = data.get("content", "").strip()
#     if not demand_content:
#         return jsonify({"error": "需求内容不能为空"}), 400
#
#     # 1. 检索知识库
#     docs = vector_store.similarity_search(demand_content, k=10) if vector_store else []
#     private_ctx = "\n".join([d.page_content for d in docs])
#
#     # 2. 实时搜索
#     web_ctx = format_search_results(search([demand_content]))
#
#     # 3. 构建提示词
#     user_prompt = f"【公司知识】{private_ctx}\n【实时搜索】{web_ctx}\n用户需求：{demand_content}"
#
#     # 4. 调用 DeepSeek
#     try:
#         response = client.chat.completions.create(
#             model=MODEL,
#             messages=[
#                 {"role": "system", "content": _make_system_prompt()},
#                 {"role": "user", "content": user_prompt}
#             ],
#             temperature=0.5,
#             max_tokens=2048,
#             top_p=0.95
#         )
#         answer = response.choices[0].message.content.strip()
#         return jsonify({"solution": answer})
#     except Exception as e:
#         return jsonify({"error": str(e)}), 500
#第2种类型
# 加一个免登录路由，只供内网调用
@app.route('/api/inner/ask', methods=['POST'])
def inner_ask():
    """
    内部系统专用，免登录，返回 JSON（流式先简化，后续再升级）
    """
    data = request.get_json(force=True)
    question = data.get("question", "").strip()
    if not question:
        return jsonify({"code": 400, "msg": "问题为空"})

    # ===== 下面直接复用你留言板现有的逻辑 =====
    # 1. 搜知识库
    docs = vector_store.similarity_search(question, k=15) if vector_store else []
    private_ctx = "\n".join([d.page_content for d in docs])
    # 2. 搜网页兜底
    web_ctx = format_search_results(search([question]))
    # 3. 构造 prompt
    user_prompt = f"【公司知识】{private_ctx}\n【实时搜索】{web_ctx}\n用户问题：{question}"

    # 4. 调 DeepSeek（非流，方便 Java 调；后续可改流）
    try:
        rsp = client.chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "system", "content": _make_system_prompt()},
                {"role": "user", "content": user_prompt}
            ],
            temperature=0.5,
            max_tokens=2048,
            top_p=0.95,
            stream=False          # 先一次性返回
        )
        answer = rsp.choices[0].message.content
        return jsonify({"code": 200, "data": answer})
    except Exception as e:
        return jsonify({"code": 500, "msg": str(e)})


# ------------------ 启动 ------------------
if __name__ == '__main__':
    print('进入 __main__')
    #运行预检函数：检查缺失的环境变量、目录是否存在，有问题提前报 warning，避免运行中途才发现
    check_configurations()
    #数据库不存在才初始化：防止重复建表或覆盖已有数据
    if not os.path.exists('chat.db'):
        print('准备创建数据库')
        with app.app_context():
            db.create_all()
            #手动推入 Flask 应用上下文：db.create_all() 必须在应用上下文里执行，否则找不到 app 配置。
            #根据已定义的 User/Message 模型自动建表：如果表已存在则跳过（不会覆盖）
            print('create_all 已执行')
            if not User.query.filter_by(username='admin').first():
            #检查是否已有 admin 账号：避免重复插入
                admin = User(username='admin', is_admin=True)
                #实例化管理员用户，is_admin=True 表示管理员权限
                admin.set_password('admin')
            #给管理员设置初始密码（明文 'admin' 被哈希后存入 password_hash 字段）
                db.session.add(admin)
            #把新用户加入数据库会话（此时还未真正写入磁盘）
                db.session.commit()
            #提交事务：真正把记录写进 chat.db，完成创建
                print('admin 用户已创建')

    app.run(host='0.0.0.0', port=5000, debug=False)

# 启动 Flask 内置服务器：
# host='0.0.0.0' → 监听所有网卡，局域网/公网都能访问；
# port=5000 → 默认端口；
# debug=False → 关闭调试模式，生产环境必须 False，防止暴露调试器与自动重载